import torch
import torchvision.transforms as transforms
from PIL import Image
import matplotlib.pyplot as plt
from pathlib import Path
from models.resnet18 import Resnet18
import utils

# 设置中文显示
plt.rcParams["font.family"] = ["SimHei"]
plt.rcParams["axes.unicode_minus"] = False  # 正确显示负号

# 模型路径和图片路径配置
MODEL_PATH = Path(__file__).parent / 'results' / 'models' / 'final_model.pth'
TEST_IMAGE_PATH = Path(__file__).parent / 'dataset' / 'valid' / 'car.jpg'
CLASSES = utils.CLASSES

# 图片预处理（针对大尺寸图片调整为等比例缩放而非裁剪）
transform = transforms.Compose([
    transforms.Resize((32, 32)),  # 将大图片等比例缩放到模型输入尺寸
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
])


def load_model(model_path):
    """加载训练好的模型"""
    model = Resnet18()
    # 加载checkpoint并提取模型参数
    checkpoint = torch.load(model_path, map_location=torch.device('cpu'))
    model.load_state_dict(checkpoint['model_state_dict'])
    model.eval()
    return model


def predict_image(model, image_path, transform):
    """预测单张图片的类别"""
    # 加载并预处理图片
    image = Image.open(image_path).convert('RGB')
    image_tensor = transform(image).unsqueeze(0)  # 添加批次维度

    # 预测
    with torch.no_grad():
        outputs = model(image_tensor)
        _, predicted = torch.max(outputs, 1)
        predicted_class = CLASSES[predicted.item()]
        confidence = torch.nn.functional.softmax(outputs, dim=1)[0][predicted.item()].item() * 100

    return image, predicted_class, confidence


def main():
    # 检查模型文件是否存在
    if not MODEL_PATH.exists():
        print(f"错误：模型文件 {MODEL_PATH} 不存在")
        return

    # 检查测试图片是否存在
    if not TEST_IMAGE_PATH.exists():
        print(f"错误：测试图片 {TEST_IMAGE_PATH} 不存在")
        return

    # 加载模型
    print(f"加载模型: {MODEL_PATH}")
    model = load_model(MODEL_PATH)

    # 预测图片
    print(f"测试图片: {TEST_IMAGE_PATH}")
    image, predicted_class, confidence = predict_image(model, TEST_IMAGE_PATH, transform)

    # 显示结果
    plt.figure(figsize=(8, 6))
    plt.imshow(image)
    plt.title(f'预测类别: {predicted_class} (置信度: {confidence:.2f}%)')
    plt.axis('off')
    plt.tight_layout()
    plt.show()

if __name__ == '__main__':
    main()